package main.java.LeetCode.ClassicalAlgorithm.Tree.Iterator;

import main.java.LeetCode.ClassicalAlgorithm.Tree.TreeIterator;
import main.java.LeetCode.CommonSolution.TreeNode;

import java.util.ArrayList;
import java.util.HashSet;

public class PreTreeIterator implements TreeIterator {
    private TreeNode cur;
    private final ArrayList<TreeNode> stack;
    private int index = 0;
    private final HashSet<TreeNode> nodeSet;

    public PreTreeIterator(TreeNode root){
        this.cur = root;
        this.stack = new ArrayList<>();
        this.nodeSet = new HashSet<>();
    }
    @Override
    public int next() {
        while (cur!=null){
            stack.add(cur);
            cur = cur.left;
        }
        TreeNode now = stack.get(index);


        if(index>=stack.size()-1) {
            for (int i = stack.size() - 1; i >= 0; i--) {
                TreeNode node = stack.get(i);
                if (node.right != null && !nodeSet.contains(node)) {
                    cur = node.right;
                    nodeSet.add(node);
                    break;
                }
            }
        }
        index++;
        return now.val;
    }

    @Override
    public boolean hasNext() {
        if(index==0&&stack.isEmpty())
            return true;
        return !(cur==null&&index>=stack.size());
    }
}
